import numpy as np
import h5py 
import cv2

# arr1 = np.random.randn(10000) 
# arr2 = np.random.randn(10000) 
# with h5py.File('test_read.hdf5', 'w') as f: 
# 	f.create_dataset('rgb', data = arr1) 
# 	f.create_dataset('depth', data = arr2) 

file="pydata.txt"

mnid = np.array([int(x.split(' ')[0]) for x in open(file).readlines()])
kp_x = np.array([float(x.split(' ')[1]) for x in open(file).readlines()])
kp_y = np.array([float(x.split(' ')[2]) for x in open(file).readlines()])
mp_x = np.array([float(x.split(' ')[3]) for x in open(file).readlines()])
mp_y = np.array([float(x.split(' ')[4]) for x in open(file).readlines()])
mp_z = np.array([float(x.split(' ')[5]) for x in open(file).readlines()])
ct_x = np.array([float(x.split(' ')[6]) for x in open(file).readlines()])
ct_y = np.array([float(x.split(' ')[7]) for x in open(file).readlines()])
ct_z = np.array([float(x.split(' ')[8]) for x in open(file).readlines()])

alldata = np.stack([mnid, kp_x, kp_y,mp_x,mp_y,mp_z,ct_x,ct_y,ct_z],axis=1)
frameWiseData = np.split(alldata, np.where(np.diff(alldata[:,0]))[0]+1)

for frame in frameWiseData:
	id = str(int(frame[0][0]))
	dist = np.linalg.norm(frame[:,3:6] - frame[:,6:9], axis=1) 
	dist = (dist*100) / np.linalg.norm(dist) # normalize to give uniform color shade
	rgb = cv2.imread("../../Downloads/nyudepthv2_basement/orbNomen2/frame"+id+".jpg")
	depth = cv2.imread("../../Downloads/nyudepthv2_basement/orbNomen/frame"+id+".jpg",0)

	# print(depth.shape) #(480, 640, 3)
	# print(rgb.shape) #(480, 640, 3)

	sparse_depth = np.zeros(rgb.shape[0:2])
	for (x, y, d) in zip(frame[:,1], frame[:,2], dist):
		sparse_depth[int(y),int(x)]=d

	# cv2.imwrite("test.jpg",sparse_depth)
	sparse_depth = cv2.flip(sparse_depth, 0) 

	with h5py.File("val/frame_" + id + ".h5", 'w') as f: 
		f.create_dataset('rgb', data = rgb) 
		f.create_dataset('depth', data = depth) 
		f.create_dataset('sparse', data = sparse_depth) 

	# break